package skewreduce.framework;

import java.util.ArrayList;
import java.util.IdentityHashMap;
import java.util.Iterator;
import java.util.NoSuchElementException;

import skewreduce.framework.physical.PhysicalOp;


public class MergeOpTree implements Iterable<MergeOpTree.Node> {
    private Node treeRoot;
    private IdentityHashMap<PhysicalOp,MergeOpTree.Node> mergeTreeMap = new IdentityHashMap<PhysicalOp,MergeOpTree.Node>();

    // for merge optimization
    public static class Node {
        PhysicalOp op;
        Node left;
        Node right;

        int traverseState;

        // rebuild every time. should post order visit

        long inputBytes;    // size from complted descendents
        int  complted;      // completed descendents so far
        int  totalComplete; // how many descendents there are

        Node(PhysicalOp op) { this.op = op; }

        public Node getLeft() { return left; }
        public Node getRight() { return right; }

        public PhysicalOp getOp() { return op; }
        public PhysicalOp getLeftOp() { return left.getOp(); }
        public PhysicalOp getRightOp() { return right.getOp(); }
        public boolean isLeaf() { return left == null && right == null; }

        public boolean isReady() {
            if ( isLeaf() )
                return op.isReady();
            else
                return left.isReady() && right.isReady();
        }

        public void setChild(Node c) {
            switch ( op.getBound().whichChild(c.getOp().getBound()) ) {
                case 0: left = c; break;
                case 1: right = c; break;
                default:
                        throw new IllegalArgumentException(c+" is not a child of "+op);
            }
        }

        public int hashCode() { return op.hashCode(); }

        public boolean equals(Object o) {
            if ( this == o ) return true;
            if ( o instanceof PhysicalOp )
                return op.equals(o);
            if ( o instanceof Node )
                return op.equals(((Node)o).op);
            return false;
        }

        public void aggregate() {
            if ( isLeaf() ) {
                // get it from children 
            } else {
            }
        }
    }

    private Node getTreeNode(PhysicalOp op) {
        Node node = mergeTreeMap.get(op);
        if ( node == null ) {
            node = new Node(op);
            mergeTreeMap.put(op,node);
        }
        return node;
    }

    void setRoot(PhysicalOp op) {
        treeRoot = getTreeNode(op);
    }

    public void addParentChild(PhysicalOp p,PhysicalOp c) {
        Node parent = getTreeNode(p);
        Node child = getTreeNode(c);

        parent.setChild(child);
    }

    /**
     * prune subtree on completion of <code>op</code>
     */
    public void remove(PhysicalOp op) {
        Node n = getTreeNode(op);
        Node l = n.getLeft();
        Node r = n.getRight();
        if ( l != null ) {
            mergeTreeMap.remove(l.getOp());
            n.left = null;
        }
        if ( r != null ) {
            mergeTreeMap.remove(r.getOp());
            n.right = null;
        }
    }

    public Iterator<MergeOpTree.Node> iterator() {
        return new PostOrderIterator(treeRoot);
    }

    // post order traversal
    // use traverseState
    // 0 : on first visit. begin left subtree
    // 1 : done left subtree. begin right subtree
    // 2 : done right subtree. done with this node.

    class PostOrderIterator implements Iterator<MergeOpTree.Node> {
        Node current;
        ArrayList<MergeOpTree.Node> stack = new ArrayList<MergeOpTree.Node>(8);

        PostOrderIterator(Node root) {
            root.traverseState = 0;
            stack.add(root);
            findNextPostOrder();
        }

        private void findNextPostOrder() {
            while ( ! stack.isEmpty() ) {
                Node top = stack.get( stack.size() - 1 );
                if ( top.isLeaf() ) return; // hooray!

                if ( top.traverseState == 0 ) {
                    // lookup left subtree
                    Node left = top.getLeft();
                    left.traverseState = 0;
                    stack.add(left);

                    ++top.traverseState;
                } else if ( top.traverseState == 1 ) {
                    // done with left subtree
                    Node right = top.getRight();
                    right.traverseState = 0;
                    stack.add(right);

                    ++top.traverseState;
                } else {
                    // it's about to leave!
                    return;
                }
            }
        }

        public boolean hasNext() {
            return ! stack.isEmpty();
        }

        public MergeOpTree.Node next() throws NoSuchElementException {
            if ( stack.isEmpty() ) {
                throw new NoSuchElementException();
            }
            Node current = stack.remove( stack.size() - 1 );
            findNextPostOrder();
            return current;
        }

        public void remove() {
            throw new UnsupportedOperationException();
        }
    }
}
